#!/usr/bin/env python3
# usage: python random_select_xyz.py [-i 输入文件] [-o 输出文件] [-n 选择的帧数]
import argparse
import random
import re
"""
随机选择指定数量的原xyz 文件中的帧
"""
def read_frames(path):
    """以纯数字行（原子总数）作为帧头进行分割"""
    with open(path, 'r', encoding='utf-8') as f:
        content = f.read()

    # 正则：以纯数字行开头，直到下一个纯数字行之前
    pattern = re.compile(r'^\d+\s*$', re.MULTILINE)
    starts = [m.start() for m in pattern.finditer(content)]
    if not starts:
        raise ValueError('未找到原子总数行，无法分割')

    frames = []
    for i, st in enumerate(starts):
        end = starts[i+1] if i+1 < len(starts) else len(content)
        frames.append(content[st:end].rstrip('\n'))

    return frames

def write_frames(frames, path):
    with open(path, 'w', encoding='utf-8') as f:
        for blk in frames:
            f.write(blk + '\n')

def main():
    parser = argparse.ArgumentParser(description='从 xyz 文件中随机选择指定数量的帧')
    # parser.add_argument('-i', '--input', type=str, default='train.xyz', help='输入 xyz 文件名（默认 input.xyz）')
    # parser.add_argument('-o', '--output', type=str, default='train_small.xyz', help='输出 xyz 文件名（默认 output.xyz）')
    parser.add_argument('-i', '--input', type=str, default='test.xyz', help='输入 xyz 文件名（默认 input.xyz）')
    parser.add_argument('-o', '--output', type=str, default='test_small.xyz', help='输出 xyz 文件名（默认 output.xyz）')
    parser.add_argument('-n', '--num_frames', type=int, default=200, help='选择的帧数（默认 800）')
    args = parser.parse_args()

    frames = read_frames(args.input)
    total_n = len(frames)
    if args.num_frames > total_n:
        import sys
        sys.exit(f'总帧数 {total_n} 小于要求的帧数 {args.num_frames}')

    random.shuffle(frames)
    selected_frames = frames[:args.num_frames]

    write_frames(selected_frames, args.output)
    print(f'完成：{args.output} = {len(selected_frames)} 帧')

if __name__ == '__main__':
    main()